%----------------------------------------------
% Formatting and file path
%----------------------------------------------

clear;
clc;
format bank; % Set formatting to display numbers without e format

%----------------------------------------------
% Read in data
%----------------------------------------------
Folder = cd; 
Folder = fullfile(Folder, '../../Data/structural/');
cd(Folder); 
load('output/Structural_Results_at_ss_assets_minus_land_of_rich.mat')

data = csvread('intermediate/w3_capital_of_stups.csv');
hhid=data(:,1);
k_w3=data(:,4);
l_w3=data(:,2);
h_w3=data(:,3);
hired_l_w3=data(:,5);
w_w3=data(:,6);
N=length(hhid); 

%----------------------------------------------
% Simplify terms
%----------------------------------------------

fk_w3 = a.*k_w3.^2 + b.*k_w3;


%----------------------------------------------
% Estimate h*, l*, h' at wave 3 capital level
%----------------------------------------------

psi_l_sim = ratio_psi_l.*psi_l;
psi_h_sim = ratio_psi_l.*psi_h;

l_opt_w3 = zeros(N,1);
h_opt_w3 = zeros(N,1);
hiredin_opt_w3 = zeros(N,1);
profit_opt_w3 = zeros(N,1);

for i=1:N
    profit = @(x) -A(i)*fk_w3(i)*(x(1) + x(3))^beta - w_w3(i)*x(2) + phi*w_w3(i)*x(3) + 0.5*(sqrt(psi_l_sim(i))*x(1) + sqrt(psi_h_sim(i))*x(2))^2;
    [argmax, max] = fmincon(profit, [0,0,0], [1,1,0], R, [], [], [0,0,0], [R,H(i),Hire_in_max]);
    l_opt_w3(i) = argmax(1);
    h_opt_w3(i) = argmax(2);
    hiredin_opt_w3(i) = argmax(3);
    profit_opt_w3(i) = -max;
end

clear max argmax

% Check if it would be optimal to liquidate capital:
for i=1:N
    profitalt = @(x) -rho*k_w3(i) - w_w3(i)*x + 0.5*psi_h_sim(i)*x^2;
    [argmax, max] = fmincon(profitalt, 0, [], [], [], [], 0, H(i));
    if -max > profit_opt_w3(i)
        l_opt_w3(i) = 0;
        hiredin_opt_w3(i) = 0;
        h_opt_w3(i) = argmax;
        profit_opt_w3(i) = -max;
    end
end

clear max argmax

% Trim :
l_opt_w3(l_opt_w3 < 1) = 0;
h_opt_w3(h_opt_w3 < 1) = 0;
hiredin_opt_w3(hiredin_opt_w3 < 1) = 0;

% Optimal cases:
case_opt_w3 = zeros(N,1);

for i=1:N
    if l_opt_w3(i) > 0 && h_opt_w3(i) > 0 && hiredin_opt_w3(i) > 0
        case_opt_w3(i) = 1;
    elseif l_opt_w3(i) > 0 && h_opt_w3(i) > 0 && hiredin_opt_w3(i) == 0
        case_opt_w3(i) = 2;
    elseif l_opt_w3(i) > 0 && h_opt_w3(i) == 0 && hiredin_opt_w3(i) > 0
        case_opt_w3(i) = 3;
    elseif l_opt_w3(i) > 0 && h_opt_w3(i) == 0 && hiredin_opt_w3(i) == 0
        case_opt_w3(i) = 4;
    elseif l_opt_w3(i) == 0 && h_opt_w3(i) > 0 && hiredin_opt_w3(i) > 0
        case_opt_w3(i) = 5;
    elseif l_opt_w3(i) == 0 && h_opt_w3(i) > 0 && hiredin_opt_w3(i) == 0
        case_opt_w3(i) = 6;
    elseif l_opt_w3(i) == 0 && h_opt_w3(i) == 0
        case_opt_w3(i) = 7;
    end
end

%----------------------------------------------
% Determine actual case in wave 3
%----------------------------------------------

case_real_w3 = zeros(N,1);

for i=1:N
    if l_w3(i) > 0 && h_w3(i) > 0 && hired_l_w3(i) > 0
        case_real_w3(i) = 1;
    elseif l_w3(i) > 0 && h_w3(i) > 0 && hired_l_w3(i) == 0
        case_real_w3(i) = 2;
    elseif l_w3(i) > 0 && h_w3(i) == 0 && hired_l_w3(i) > 0
        case_real_w3(i) = 3;
    elseif l_w3(i) > 0 && h_w3(i) == 0 && hired_l_w3(i) == 0
        case_real_w3(i) = 4;
    elseif l_w3(i) == 0 && h_w3(i) > 0 && hired_l_w3(i) > 0
        case_real_w3(i) = 5;
    elseif l_w3(i) == 0 && h_w3(i) > 0 && hired_l_w3(i) == 0
        case_real_w3(i) = 6;
    elseif l_w3(i) == 0 && h_w3(i) == 0
        case_real_w3(i) = 7;
    end
end


%----------------------------------------------
% Estimate value of misallocation at wave 3
%----------------------------------------------

misallocation_w3 = zeros(N,1);

for i=1:N
    if case_opt(i) == 5 || case_opt(i) == 6   % Misallocation is zero if case 5 or 6 dominates at high SS
        misallocation_w3(i) = 0;
    else
        misallocation_w3(i) = profit_opt(i) - profit_opt_w3(i);
    end
end

misallocation_w3(misallocation_w3 < 0) = 0;

% Top-code top 5% outliers
misallocation_95_pcile_w3 = prctile(misallocation_w3,95);
misallocation_topcoded_w3 = misallocation_w3;
misallocation_topcoded_w3(misallocation_topcoded_w3>misallocation_95_pcile_w3) = misallocation_95_pcile_w3;

total_misallocation_w3=sum(misallocation_w3);
total_misallocation_topcoded_w3=sum(misallocation_topcoded_w3);

% Calculate total number of individuals misallocated

num_misallocated_w3 = sum(misallocation_w3(:) ~= 0);


%----------------------------------------------
% Export results to Excel
%----------------------------------------------

T=table(hhid,case_opt_w3,case_real_w3,l_opt_w3,l_w3,h_opt_w3,h_w3,hiredin_opt_w3,hired_l_w3,k_w3,psi_h,psi_l,A,k_bl_post_trans);
filename='output/w3_test_of_model.xlsx';
writetable(T,filename);

save('output/w3_test_results')
